import pickle
from argparse import ArgumentParser
from ibm_model import IBMModel1


def load_sent_pairs(source_file, target_file):
    pairs = []

    source_file = open(source_file, 'r')
    target_file = open(target_file, 'r')

    for src_line, tgt_line in zip(source_file, target_file):
        src_line = [int(i) for i in src_line.strip().split()]
        tgt_line = [int(i) for i in tgt_line.strip().split()]
        pair = (src_line, tgt_line)
        pairs.append(pair)

    source_file.close()
    target_file.close()

    return pairs


def main(option):
    src_vocab = pickle.load(open(option.source_vocab, 'rb'))
    tgt_vocab = pickle.load(open(option.target_vocab, 'rb'))

    pairs = load_sent_pairs(option.source_file, option.target_file)

    lf = len(src_vocab['word2id'])
    le = len(tgt_vocab['word2id'])

    model = IBMModel1(le, lf)

    das = src_vocab['word2id']['das']
    buch = src_vocab['word2id']['Buch']
    haus = src_vocab['word2id']['Haus']
    ein = src_vocab['word2id']['ein']

    the = tgt_vocab['word2id']['the']
    book = tgt_vocab['word2id']['book']
    house = tgt_vocab['word2id']['house']
    a = tgt_vocab['word2id']['a']

    for i in range(5):
        model.train(pairs)
        print('Epoch ' + str(i + 1))
        print('the      das     ' + str(model.t[the][das]))
        print('book     das     ' + str(model.t[book][das]))
        print('house    das     ' + str(model.t[house][das]))
        print('the      Buch    ' + str(model.t[the][buch]))
        print('book     Buch    ' + str(model.t[book][buch]))
        print('a        Buch    ' + str(model.t[a][buch]))
        print('book     ein     ' + str(model.t[book][ein]))
        print('a        ein     ' + str(model.t[a][ein]))
        print('the      Haus    ' + str(model.t[the][haus]))
        print('house    Haus    ' + str(model.t[house][haus]))


if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument('--source_file', type=str, default='output/test.gr.txt')
    parser.add_argument('--target_file', type=str, default='output/test.en.txt')
    parser.add_argument('--source_vocab', type=str, default='output/vocab_src.pkl')
    parser.add_argument('--target_vocab', type=str, default='output/vocab_tgt.pkl')
    option = parser.parse_args()

    main(option)
